
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/err.h>
#include "testAESKeys.h"
#include "../utils/StringUtils.h"

void tests::testAESKeys(apr_pool_t *mp)
{
   // Can generate an AES symmetric key by using the constructor as:
   AESSymmetricKey *aesky = new AESSymmetricKey(256, "newAES");
   int iters = 10;
   testAESBlockCryption(aesky, iters, 1000000);
   testAESPasswordCryption(aesky, iters);
   testAESPipeCryption(aesky, iters, 50000, 1000000, mp);

   delete aesky;
   ERR_remove_state(0);
}

void tests::testAESBlockCryption(AESSymmetricKey *aesky, int iterations, int max_length)
{
   for (int i = 0; i < iterations; i++)
   {
      // Random length much more than a single key
      int randlen = rand() % max_length;
      // Create a block of (pseudo-)random data
      char *pseudorandom = (char *)malloc(sizeof(char)*randlen);
      for (int j = 0; j < randlen; j++)
         pseudorandom[j] = (rand() & 0xFF0) >> 8;
      // Encrypt the random data
      apr_size_t new_coded_size = randlen;
      apr_size_t new_plain_size = 0;
      AESEncrypt *aes_enc = new AESEncrypt(*aesky);
      char *encrypted_data = aes_enc->processDataBlock(pseudorandom, randlen, &new_coded_size);
      delete aes_enc;
      AESDecrypt *aes_dec = new AESDecrypt(*aesky);
      char *new_plain_data = aes_dec->processDataBlock(encrypted_data, new_coded_size, &new_plain_size);
      delete aes_dec;
      if (new_plain_size != randlen)
      {
         errprint("error: original data size was %" APR_SIZE_T_FMT ", new data size is %" APR_SIZE_T_FMT, (apr_size_t)randlen, new_plain_size);
      } else if (memcmp(pseudorandom, new_plain_data, sizeof(char)*randlen) != 0)
      {
         errprint("error: original data does not match new data");
         //printf("old data:\n");
         //StringUtils::printBinary(pseudorandom, randlen);
         //printf("new data:\n");
         //StringUtils::printBinary(new_plain_data, randlen);
      } else
      {
         printf("AES BlockCryption test %d passed - data length %d\n", i, randlen);
      }
      free(pseudorandom);
      free(encrypted_data);
      free(new_plain_data);
   }
}

void tests::testAESPasswordCryption(AESSymmetricKey *aesky, int iterations)
{
   for (int i = 0; i < iterations; i++)
   {
      char password_buffer[256];
      char *exporttext = NULL;
      sprintf(password_buffer, "password%d", rand());
      aesky->exportKeyToASCIIWithPassword(&exporttext, password_buffer);
      AESSymmetricKey *aesky2 = new AESSymmetricKey("newAES");
      aesky2->importKeyFromASCIIWithPassword(exporttext, password_buffer, "newAES");
      int compare_key = aesky2->compareKeyEqual(aesky);
      if (compare_key)
      {
         printf("AES PasswordCryption test %d passed - password '%s'\n", i, password_buffer);
      } else
      {
         errprint("failed: encrypted key was '%s'", exporttext);
         errprint("original key base64 '%s'", aesky->getKeyData());
         errprint("final key base64 '%s'", aesky2->getKeyData());
      }
      free(exporttext);
      delete aesky2;
   }
}

void tests::testAESPipeCryption(AESSymmetricKey *aesky, int iterations, int max_block, int max_length, apr_pool_t *mp)
{
   for (int i = 0; i < iterations; i++)
   {
      // Random length much more than a single key
      int randlen = rand() % max_length;
      // Create a block of (pseudo-)random data
      char *pseudorandom = (char *)malloc(sizeof(char)*randlen);
      for (int j = 0; j < randlen; j++)
         pseudorandom[j] = (rand() & 0xFF0) >> 8;
      // Encrypt the random data
      apr_size_t new_coded_size = randlen;
      APRBucketBrigade p_input(mp);
      APRBucketBrigade coded(mp);
      AESEncrypt *aes_enc = new AESEncrypt(p_input, coded, *aesky);
      aes_enc->processPipeBegin();
      int nidx = 0;
      while (nidx < randlen)
      {
         int randblocklen = rand() % max_block;
         randblocklen = ((randblocklen + nidx) > randlen) ? (randlen - nidx) : randblocklen;
         int randblocklen_added = p_input.writeToStream(&(pseudorandom[nidx]), randblocklen);
         nidx += randblocklen_added;
         aes_enc->processPipeStep();
      }
      aes_enc->processPipeEnd();
      delete aes_enc;
      APRBucketBrigade coded2(mp);
      APRBucketBrigade p_output(mp);
      AESDecrypt *aes_dec = new AESDecrypt(coded2, p_output, *aesky);
      aes_dec->processPipeBegin();
      nidx = 0;
      int total_length = coded.realLength();
      while (nidx < total_length)
      {
         int randblocklen = rand() % max_block;
         randblocklen = ((randblocklen + nidx) > total_length) ? (total_length - nidx) : randblocklen;
         char *working_buffer = (char *)malloc(sizeof(char)*randblocklen);
         int block_read = coded.readFromStream(working_buffer, randblocklen);
         if (block_read != randblocklen)
         {
            errprint("read wrong amount from bucket (expected %d, got %d)", randblocklen, block_read);
         } else
            coded.removeFromStream(randblocklen);
         int block_write = coded2.writeToStream(working_buffer, randblocklen);
         if (block_write != randblocklen)
         {
            errprint("wrote wrong amount to bucket (expected %d, wrote %d)", randblocklen, block_write);
         }
         nidx += randblocklen;
         aes_dec->processPipeStep();
         free(working_buffer);
      }
      aes_dec->processPipeEnd();
      delete aes_dec;
      size_t new_plain_size = 0;
      char *new_plain_data = NULL;
      p_output.bucketBrigadeToBuffer(&new_plain_data, &new_plain_size);
      if (new_plain_size != randlen)
      {
         errprint("error: original data size was %" APR_SIZE_T_FMT ", new data size is %" APR_SIZE_T_FMT, (apr_size_t)randlen, new_plain_size);
         printf("old data:\n");
         StringUtils::printBinary(pseudorandom, randlen);
         printf("new data:\n");
         StringUtils::printBinary(new_plain_data, randlen);
      } else if (memcmp(pseudorandom, new_plain_data, sizeof(char)*randlen) != 0)
      {
         errprint("error: original data does not match new data");
         //printf("old data:\n");
         //StringUtils::printBinary(pseudorandom, randlen);
         //printf("new data:\n");
         //StringUtils::printBinary(new_plain_data, randlen);
      } else
      {
         printf("AES PipeCryption test %d passed - data length %d\n", i, randlen);
      }
      free(pseudorandom);
      free(new_plain_data);
   }
}
